package com.github.dongliusuo.sshtunneldemo;

import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.io.IOException;

/**
 * 建立SSH隧道
 * 要确保远程服务器的 SSH 配置文件 /etc/ssh/sshd_config，
 * 确保以下选项被启用：
 *          AllowTcpForwarding yes
 * 如果没启用，修改配置后重启sshd，
 * 命令: sudo systemctl restart sshd
 */
public class SshTunnel implements Closeable {

    private static final Logger log = LoggerFactory.getLogger(SshTunnel.class);

    private final SshTunnelProperties properties;

    private Session session;

    public SshTunnel(SshTunnelProperties properties) {
        this.properties = properties;
        log.info("### Init SSH Tunnel, config={}", properties);
        try {
            JSch jSch = new JSch();
            if (properties.getPrivateKeyPath()!= null && !properties.getPrivateKeyPath().isBlank()) {
                // 用私钥
                jSch.addIdentity(properties.getPrivateKeyPath(), properties.getSshPassword());
                this.session = jSch.getSession(properties.getSshUser(), properties.getSshHost(), properties.getSshPort());
            } else {
                // 用密码
                this.session = jSch.getSession(properties.getSshUser(), properties.getSshHost(), properties.getSshPort());
                this.session.setPassword(properties.getSshPassword());
            }
            // 避免首次连接时出现是否信任主机的提示
            this.session.setConfig("StrictHostKeyChecking", "no");
        } catch (JSchException e) {
            throw new RuntimeException(e);
        }
    }

    public void open() {
        try {
            log.info("### SSH connecting...");
            this.session.connect(this.properties.getSshConnectTimeout());
            log.info("### SSH connected! Starting PortForwardingL...");
            this.session.setPortForwardingL(this.properties.getLocalPort(), this.properties.getRemoteHost(), this.properties.getRemotePort());
            log.info("### SSH Tunnel open: localhost:{} -> {}:{}",
                    this.properties.getLocalPort(),
                    "localhost".equals(this.properties.getRemoteHost()) ? this.properties.getSshHost() : this.properties.getRemoteHost(),
                    this.properties.getRemotePort()
            );
        } catch (JSchException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void close() throws IOException {
        log.info("### SSH start disconnect...");
        if (session != null) {
            session.disconnect();
            log.info("### SSH disconnected!");
        }
    }
}
